#include "utils.h"
#include <base/utility.h>
#include <base/geometry.h>
#include <base/macro.h>
#include <base/log.h>

#include <Eigen/Dense>
#include <Eigen/Core>

using namespace SPM;


bool utils::Matching(const DlFrame &dl,
			const hdmap_frame& hdmap_f,
			const Camera &cam,
			const Pose &pose,
            std::vector<std::pair<Eigen::Vector2d, Eigen::Vector3d>>& match_results)
{
	 Eigen::Matrix3d bRw = pose.orientation.toRotationMatrix();
     Eigen::Vector3d bTw = pose.position;
     Eigen::Matrix3d cRb = cam.mount.rotation.matrix();
     Eigen::Vector3d cTb = cam.mount.translation;

     std::vector<DlSign> dlSigns = dl.signs;
	 std::vector<std::pair<DlSign, hdmap_sign>> sign_matches;
     for (hdmap_sign sign : hdmap_f.hdmap_signs)
     {
         std::vector<Eigen::Vector2d> sign_corners;
         for (std::size_t i = 0; i < sign.geometry.size(); i++)
         {
             Eigen::Vector2d uv;
             Eigen::Vector3d pcd = cRb.transpose() * (bRw.transpose() * (sign.geometry[i] - bTw) - cTb);
             if (Utility::Project3d_2d(pcd, uv, cam.K))
             {
                 sign_corners.emplace_back(uv.x(), uv.y());
             }
         }
         envelope2D sign_env = envelope2D::GetEnvelope(sign_corners);

         int matchIdx = -1;
         double max_overlop = 0.0;
         for(std::size_t i = 0; i < dlSigns.size(); i++)
         {
             envelope2D bounding(dlSigns[i].bBox);
             double iou = envelope2D::IOU(bounding, sign_env);
             if(iou > 0.3 && iou > max_overlop)
             {
                 max_overlop = iou;
                 matchIdx = i;
             }
         }
         if(matchIdx != -1)
         {
             sign_matches.emplace_back(dlSigns[matchIdx], sign);
             dlSigns.erase(dlSigns.begin() + matchIdx);
         }
     }

	 for(auto match_pair : sign_matches)
	 {
		 envelope2D bBox = match_pair.first.bBox;
		 Eigen::Vector2d min_env, max_env;
		 bBox.GetRange(min_env, max_env);
		 std::vector<Eigen::Vector3d> body_vetex(4);
		 for(std::size_t i = 0; i < 4; i++)
		 {
		 	body_vetex[i] = bRw.transpose() * (match_pair.second.geometry[i] - bTw);
		 }
		 match_results.emplace_back(Eigen::Vector2d(max_env.x(), max_env.y()), body_vetex[0]);
		 match_results.emplace_back(Eigen::Vector2d(min_env.x(), max_env.y()), body_vetex[1]);
		 match_results.emplace_back(Eigen::Vector2d(min_env.x(), min_env.y()), body_vetex[2]);
		 match_results.emplace_back(Eigen::Vector2d(max_env.x(), min_env.y()), body_vetex[3]);
	 }
	 return true;
}


bool utils::Translate3d_2d(const std::vector<std::pair<Eigen::Vector2d, Eigen::Vector3d>>& pole_sign_matches,  
            const Camera& camera,
            Eigen::Matrix3d& R,
            Eigen::Vector3d& T, 
            double& rmse)
{
    if (pole_sign_matches.empty())
	{
		LOG("match_points is empty, translate failed!") << std::endl;
		return false;
	}

	Eigen::Matrix3d cRb = camera.mount.rotation.matrix();
	Eigen::Vector3d cTb = camera.mount.translation;

	int N = pole_sign_matches.size();
	std::vector<Eigen::Vector3d> q1(N);
    std::vector<Eigen::Vector2d> q2(N);
	for (int i = 0; i < N; i++)
	{
		q1[i] = pole_sign_matches[i].second;
		q2[i] = pole_sign_matches[i].first;
	}

	auto robustWeightCauchy = [=](double norm_res) -> double {
		return std::sqrt(1.0 / (1.0 + norm_res * norm_res));
	};

	double kappa = .0, t1 = .0;
	Eigen::MatrixXd A = Eigen::MatrixXd::Zero(2 * N, 2);
	Eigen::MatrixXd B = Eigen::MatrixXd::Zero(2 * N, 1);
	Eigen::MatrixXd W = Eigen::MatrixXd::Identity(2 * N, 2 * N);

	int iter = 0;
	while (iter < 20)
	{
		for (int i = 0; i < N; i++)
		{
			Eigen::Vector3d trans = geometry::RotationAroundZ(kappa) * q1[i] + Eigen::Vector3d(t1, 0, 0);
			Eigen::Vector3d cam_X = cRb.transpose() * (trans - cTb);
			  if(cam_X[2] == 0.0)
                continue;

            Eigen::Vector3d cam2D = camera.K * cam_X;
            Eigen::Vector2d uv(cam2D[0]/ cam2D[2], cam2D[1]/ cam2D[2]);
			double err_i_norm = (uv - q2[i]).norm();

			double x = trans.x();
			double y = trans.y();
			double z = trans.z();

			Eigen::MatrixXd c2uv = Eigen::MatrixXd::Zero(2, 3);
			c2uv << camera.fx() * (1 / cam_X[2]), 0.0, -camera.fx() * (cam_X[0] / (cam_X[2] * cam_X[2])),
				0.0, camera.fy() * (1 / cam_X[2]), -camera.fy() * (cam_X[1] / (cam_X[2] * cam_X[2]));

			Eigen::MatrixXd jacobian = Eigen::MatrixXd::Zero(3, 2);
			jacobian << -y, 1.0,
				x - t1, 0.0,
				0.0, 0.0;
			
			A.block(i * 2, 0, 2, 2) = c2uv * cRb.transpose() * jacobian;
			B.block(i * 2, 0, 2, 1) = q2[i] - uv;

			q1[i] = trans;
		}
		Eigen::Vector2d delta = (A.transpose() * W * A).ldlt().solve(A.transpose() * W *  B);
		double err = delta.norm();
		if (delta.norm() < EPS)
		{
			break;
		}
		kappa += delta(0);
		t1 += delta(1);
		iter++;
	}
	R = geometry::RotationAroundZ(kappa);
	T = Eigen::Vector3d(t1, 0, 0);
	return true;
}

bool utils::ComputeRT(const std::vector<std::pair<DlSign, hdmap_sign>>& track_results,
			const Pose &pose,
            const Camera &cam, 
            Eigen::Matrix3d& R,
            Eigen::Vector3d& T, 
            double& rmse,
            double tolerance,
			std::vector<std::pair<Eigen::Vector2d, Eigen::Vector3d>>& match_pairs)
{
	std::vector<std::pair<Eigen::Vector2d, Eigen::Vector3d>> match_results;
	Eigen::Matrix3d bRw = pose.orientation.toRotationMatrix();
    Eigen::Vector3d bTw = pose.position;

	for(auto match_pair : track_results)
	 {
		 envelope2D bBox = match_pair.first.calibrated_bbox;
         std::vector<Eigen::Vector3d> body_vetex(4);
		 for(std::size_t i = 0; i < 4; i++)
		 {
		 	body_vetex[i] = bRw.transpose() * (match_pair.second.geometry[i] - bTw);
		 }
		 Eigen::Vector2d min_env, max_env;
		 bBox.GetRange(min_env, max_env);
		 match_results.emplace_back(Eigen::Vector2d(max_env.x(), max_env.y()), body_vetex[0]);
		 match_results.emplace_back(Eigen::Vector2d(min_env.x(), max_env.y()), body_vetex[1]);
		 match_results.emplace_back(Eigen::Vector2d(min_env.x(), min_env.y()), body_vetex[2]);
		 match_results.emplace_back(Eigen::Vector2d(max_env.x(), min_env.y()), body_vetex[3]);
	 }

	if(Translate3d_2d(match_results, cam, R, T, rmse))
	{
		match_pairs = match_results;
		return true;
	}
	return false;
}